package com.ibm.cps.spark.streaming

import org.apache.spark.{HashPartitioner, SparkConf}
import org.apache.spark.streaming.{Seconds, StreamingContext}

/**
 * Created by telekinesis on 4/6/15.
 */
object StatefulWordCount {
  def main(args: Array[String]) {
    val updateFunc = (values: Seq[Int], state: Option[Int]) => {
      val currentCount = values.sum
      val previousCount = state.getOrElse(0)
      Some(currentCount + previousCount)
    }
    val newUpdateFunc = (iterator: Iterator[(String, Seq[Int], Option[Int])]) => {
      iterator.flatMap(t => updateFunc(t._2, t._3).map(s => (t._1, s)))
    }
    val sparkConf = new SparkConf().setAppName("StatefulNetworkWordCount")
      .set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
      .setMaster("local[2]")
      .set("spark.local.dir", "./spark-local")
      .set("spark.eventLog.enabled", "true")
      .set("spark.eventLog.dir", "./spark-events")
      .set("spark.shuffle.manager", "SORT")

    // Create the context with a 1 second batch size
    val ssc = new StreamingContext(sparkConf, Seconds(5))
    ssc.checkpoint(".")
    // Initial RDD input to updateStateByKey
//    val initialRDD = ssc.sparkContext.parallelize(List(("hello", 1), ("world", 1)))
    // Create a ReceiverInputDStream on target ip:port and count the
    // words in input stream of \n delimited test (eg. generated by 'nc')
    val lines = ssc.socketTextStream("localhost" , 9999)
    val words = lines.flatMap(_.split(" "))
    val wordDstream = words.map((_, 1))
    // Update the cumulative count using updateStateByKey
    // This will give a Dstream made of state (which is the cumulative count of the words)

    val stateDstream = wordDstream.updateStateByKey[Int](updateFunc)
//    val stateDstream = wordDstream.updateStateByKey[Int](newUpdateFunc,
//      new HashPartitioner (ssc.sparkContext.defaultParallelism), true)
    stateDstream.print()
    ssc.start()
    ssc.awaitTermination()
  }
}
